from PIL import Image
import numpy as np

import os
import os.path
from paddle.fluid.io import Dataset

def has_file_allowed_extension(filename, extensions):
    ''' Checks if a file is an allowed extension '''

    filename_lower = filename.lower()
    return any(filename_lower.endswith(ext) for ext in extensions)

def find_classes(dir):
    classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
    classes.sort()
    class_to_idx = {classes[i]: i for i in range(len(classes))}
    return classes, class_to_idx

def make_dataset(dir, extensions):
    images = []
    for root, _, fnames in sorted(os.walk(dir)):
        for fname in sorted(fnames):
            if has_file_allowed_extension(fname, extensions):
                path = os.path.join(root, fname)
                item = (path, 0)
                images.append(item)

    return images

class DatasetFolder(Dataset):
    def __init__(self, root, loader, extensions, transform=None, target_transform=None):
        samples = make_dataset(root, extensions)
        if len(samples) == 0:
            raise(RuntimeError("Found 0 files in subfolders of: " + root + "\n"
                               "Supported extensions are: " + ",".join(extensions)))

        self.root = root
        self.loader = loader
        self.extensions = extensions
        self.samples = samples

        self.transform = transform
        self.target_transform = target_transform

    def __getitem__(self, index):

        """
        :param index (int): Index
        :return tuple: (sample, target) where target is class_index of the target class.
        """

        path, target = self.samples[index]
        sample = self.loader(path)
        sample = np.asarray(sample).astype('float32')
        #sample = sample / 255.0 
        sample = sample / 127.5 -1.0   #to replace paddlex.cls.normalize
        # print("sample: ",sample)
        # print(sample.shape,np.max(sample),np.min(sample))
        if self.transform is not None:
            # print("shape of sample: ",sample.shape)
            sample = self.transform(sample)[0]
            sample = sample.transpose((2, 0, 1))
        if self.target_transform is not None:
            target = self.target_transform(target)[0]
            target=target.transpose((2, 0, 1))
        # print("sample after trans: ",sample)
        # print(sample.shape,np.max(sample),np.min(sample))

        return sample, target

    def __len__(self):
        return len(self.samples)

    def __repr__(self):
        fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
        fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
        fmt_str += '    Root Location: {}\n'.format(self.root)
        tmp = '    Transforms (if any): '
        fmt_str += '{0}{1}\n'.format(tmp, self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
        tmp = '    Target Transforms (if any): '
        fmt_str += '{0}{1}'.format(tmp, self.target_transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
        return fmt_str

IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif']

def pil_loader(path):
    with open(path, "rb") as f:
        img = Image.open(f)
        return img.convert("RGB")

def default_loader(path):
    return pil_loader(path)

class ImageFolder(DatasetFolder):
    def __init__(self, root, tranform=None, target_transform=None, loader=default_loader):
        super(ImageFolder, self).__init__(root, loader, IMG_EXTENSIONS,
                                          transform=tranform,
                                          target_transform=target_transform)
        self.imgs = self.samples




